module Main where

import Debug.QuickCheck
import System.Random
import Sample
import Data.Maybe
import Data.List

-----------------------------------------------------------------------

data Expr
  = Num Integer
  | Add Expr Expr
  | Mul Expr Expr
  | Var Name
 deriving ( Eq )

type Name = String

-----------------------------------------------------------------------

showExpr :: Expr -> String
showExpr (Num n)   = show n
showExpr (Add a b) = showExpr a ++ "+" ++ showExpr b
showExpr (Mul a b) = showFactor a ++ "*" ++ showFactor b
showExpr (Var x)   = x

showFactor :: Expr -> String
showFactor (Add a b) = "(" ++ showExpr (Add a b) ++")"
showFactor e         = showExpr e

instance Show Expr where
  show = showExpr

-----------------------------------------------------------------------

type Env = [(String,Integer)]

eval :: Env -> Expr -> Integer
eval env (Num n)   = n
eval env (Add a b) = eval env a + eval env b
eval env (Mul a b) = eval env a * eval env b
eval env (Var x)   = fromJust (lookup x env)

-----------------------------------------------------------------------

instance Arbitrary Expr where
  arbitrary = sized arbExpr

arbExpr :: Int -> Gen Expr
arbExpr s =
  frequency [ (1, do n <- arbitrary
                     return (Num n))
            , (s, do a <- arbExpr s'
                     b <- arbExpr s'
                     return (Add a b))
            , (s, do a <- arbExpr s'
                     b <- arbExpr s'
                     return (Mul a b))
            , (1, do x <- oneof [ return s | s <- ["x","y","z"] ]
                     return (Var x))
            ]
 where
  s' = s `div` 2
  
-----------------------------------------------------------------------

vars :: Expr -> [Name]
vars (Num n)   = []
vars (Add a b) = vars a `union` vars b
vars (Mul a b) = vars a `union` vars b
vars (Var y)   = [y]

-----------------------------------------------------------------------

diff :: Expr -> Name -> Expr
diff (Num n)   x = Num 0
diff (Add a b) x = Add (diff a x) (diff b x)
diff (Mul a b) x = Add (Mul a (diff b x)) (Mul b (diff a x))
diff (Var y)   x
  | x == y       = Num 1
  | otherwise    = Num 0

-----------------------------------------------------------------------
-- ANSWER EXERCISE A:

prop_DiffVars e =
  and [ y `elem` xs | y <- ys ]
 where
  xs = vars e
  ys = vars (diff e "x")

prop_DiffVars_False e =
  and [ x `elem` ys | x <- xs ]
 where
  xs = vars e
  ys = vars (diff e "x")

-----------------------------------------------------------------------
-- SIMPLE ANSWER EXERCISE B:

{-
simplify :: Expr -> Expr
simplify e | null (vars e) =
  Num (eval [] e)

simplify (Add (Add a b) c) =
  simplify (Add a (Add b c))

simplify (Add a b)
  | a' == Num 0 = b'
  | b' == Num 0 = a'
  | a' == b'    = simplify (Mul (Num 2) a')
  | otherwise   = Add a' b'
 where
  a' = simplify a
  b' = simplify b

simplify (Mul (Mul a b) c) =
  simplify (Mul a (Mul b c))

simplify (Mul a b)
  | a' == Num 0 = Num 0
  | b' == Num 0 = Num 0
  | a' == Num 1 = b'
  | b' == Num 1 = a'
  | otherwise   = Mul a' b'
 where
  a' = simplify a
  b' = simplify b

simplify e = e
-}

-- still, lots of improvements can be made to the above

-----------------------------------------------------------------------
-- COMPLICATED ANSWER EXERCISE B:
-- (this was probably too complicated for you to do during the
-- exercises. Have a look at the answer anyway.)

type Poly = [(Integer,[Name])] -- sum of products

{-
An example of a polynomial is:

  [(1,["x"]), (3,["x","y"]), (2,["x","x"]), (5,["y"])]

which represents the expression:

  x + 3*x*y + 2*x*x + 5*y
-}

-- poly converts an expression into a polynomial.
--   example: poly (Add (Var "x") (Var "y")) = [(1,["x"]),(1,["y"])]
poly :: Expr -> Poly
poly (Num 0)   = []
poly (Num n)   = [(n,[])]
poly (Var x)   = [(1,[x])]
poly (Add a b) = poly a ++ poly b
poly (Mul a b) = [ (ca*cb,xa++xb) | (ca,xa) <- poly a, (cb,xb) <- poly b ]

-- simpPoly simplifies a polynomial
--   example: simpPoly [(1,["x"]),(3,["x"])] = [(4,["x"])]
simpPoly :: Poly -> Poly
simpPoly p =
  [ ( sum [ n | (n,x') <- p', x == x' ]
    , x
    )
  | x <- Data.List.sort xs
  ]
 where
  p' = [ (n,Data.List.sort x) | (n,x) <- p ]
  xs = nub [ x | (_,x) <- p' ]

-- expr converts a (simplified) polynomial into an expression
--   example: expr [(1,["x"]),(3,["y"])] = Add (Var "x") (Mul (Num "3") 
expr :: Poly -> Expr
expr p =
  add
  [ n `mul` x
  | (n,x) <- p
  , n /= 0
  ]
 where
  n `mul` [] = Num n
  1 `mul` x  = muls x
  n `mul` x  = Num n `Mul` muls x

  muls []     = Num 1
  muls [a]    = Var a
  muls (a:as) = Var a `Mul` muls as

  add []     = Num 0
  add [a]    = a
  add (a:as) = a `Add` add as

simplify :: Expr -> Expr
simplify e = expr (simpPoly (poly e))

-----------------------------------------------------------------------

data XyzEnv = Xyz Env
 deriving ( Show )
 
instance Arbitrary XyzEnv where
  arbitrary =
    do x <- arbitrary
       y <- arbitrary
       z <- arbitrary
       return (Xyz [("x",x),("y",y),("z",z)])

prop_SimplifyCorrect e (Xyz env) =
  eval env e == eval env (simplify e)

-----------------------------------------------------------------------
-- ANSWER EXERCISE C:

prop_SimplifyNoJunk e =
  noJunk (simplify e)
 where
  -- just an example; can be extended at will!
  noJunk (Add a b) = not (isNum a && isNum b)
                  && noJunk a && noJunk b
  noJunk (Mul a b) = not (isNum a && isNum b)
                  && noJunk a && noJunk b
  noJunk _         = True

  isNum (Num _) = True
  isNum _       = False

-----------------------------------------------------------------------
-- ANSWER EXERCISE D:

prop_SimplifyDiff e =
  simplify (diff e "x") == simplify (diff (simplify e) "x")

-- It does not hold for the simple simplify function; the order of things
-- changes.
-- It could be made to hold for a perfect simplification function.
-- It holds for the more complicated simplification function.

-----------------------------------------------------------------------

